import random
import time
import sys
import onnx
import pandas
import xlwt
import numpy
import networkx as nx
import matplotlib.pyplot as plt
from collections import Counter


def inload(nosuffixname):
        model = onnx.load(nosuffixname+ '.onnx')
        model_bytes = model.SerializeToString();
        model_str = str(model)
        with open(nosuffixname + '.txt', 'w') as f:
            f.write(model_str)
        f.close()
    #输出input.xls
def toxls_input(nosuffixname):
        # 把主要信息提取，如节点名称，节点类型
        def research(nosuffixname):
            with open('input.txt', 'w') as out:
                key = ('op_type')
                key2 = ('input')
                f = open(nosuffixname + '.txt', 'r')
                l = f.readlines()
                for i in range(len(l)):
                    if key in l[i]:
                        out.write(l[i - 1])
                        out.write(l[i])
                    if key2 in l[i]:
                        out.write(l[i])
            out.close()

        # 删除常数变量的input
        def deleteConstant():
            with open('deleteConstant.txt', 'w') as input:
                keyword2 = ('op_type: "Constant"')
                keyword = ('name')
                f = open('input.txt', 'r')
                l = f.readlines()
                for i in range(len(l)):
                    if i + 1 <= len(l) - 1:
                        if keyword2 in l[i] or keyword2 in l[i + 1]:
                            pass
                        else:
                            input.write(l[i])
                    else:
                        pass
                input.close()

        # 删除非节点数据
        def drop(words: list):
            f1 = open('deleteConstant.txt')
            f2 = open('input_delete.txt', 'w')
            s = pandas.Series(f1.readlines())
            flag = numpy.array([True] * len(s))
            for word in words:
                flag = flag & ~s.str.contains(word)
            s = s[flag].tolist()
            f2.writelines(s)
            f1.close()
            f2.close()
            #把多个输入流合在一起
            temp = ''
            with open('input_emerge.txt', 'w') as input:
                keyword = ('input')
                f = open('input_delete.txt', 'r')
                flag = True
                l = f.readlines()
                for i in range(len(l)):
                    if i + 1 <= len(l):
                        if keyword in l[i]:
                            if keyword in l[i + 1]:
                                flag = False
                                strinput = l[i].split('"')[1]
                                temp = (temp + strinput + ';')
                                # print(temp)
                                input.write('')
                            elif keyword not in l[i + 1] and flag == False:
                                strline = l[i].split('"')[1]
                                temp = ('input:\"' + temp + strline + '\"' + '\n')
                                input.write(temp)
                                temp = ''
                            else:
                                input.write(l[i])
                        else:
                            input.write(l[i])
                input.close()
            i = 1
            # 对节点信息分类好，把每一个节点的信息分在同一行
            with open('input_outdelete.txt', 'w') as f:
                with open('input_emerge.txt', 'r') as fp:
                    for line in fp:
                        n = i % 3
                        if n != 0:
                            line = line.strip('\n')
                            f.write(line)
                        else:
                            f.write(line)
                        i = i + 1

        # 保存数据到CSV
        def xls_write_input(nosuffixname):
            book = xlwt.Workbook(encoding='utf-8', style_compression=0)
            sheet = book.add_sheet('input', cell_overwrite_ok=True)
            sheet.write(0, 0, 'id')
            sheet.write(0, 1, 'name')
            sheet.write(0, 2, 'type')
            sheet.write(0, 3, 'input')
            n = 1
            count = len(open('input_outdelete.txt', 'r+').readlines())
            with open('input_outdelete.txt', 'r+') as f:
                for text in f.readlines():
                    num = n
                    # name = text.split('"')[3]
                    # ran_str = ''.join(random.sample(['z','y','x','w','v','u','t','s','r','q','p','o','n','m','l','k','j','i','h','g','f','e','d','c','b','a'], 5))
                    name = str(n)
                    type = text.split('"')[5]
                    output = text.split('"')[1]
                    sheet.write(n, 0, num)
                    sheet.write(n, 1, name.split('"')[0])
                    sheet.write(n, 2, type.split('"')[0])
                    sheet.write(n, 3, output.split('"')[0])
                    n = n + 1
            book.save(nosuffixname + '_input.xls')
        research(nosuffixname)
        deleteConstant()
        drop(["node", "Constant", "{", "}"])
        xls_write_input(nosuffixname)
        print('input success')


    # 输出output.xls
def toxls_output(nosuffixname):
        # 把主要信息提取，如节点名称，节点类型
        def research(nosuffixname):
            with open('output.txt', 'w') as out:
                key = ('op_type')
                key2 = ('output')
                f = open(nosuffixname + '.txt', 'r')
                l = f.readlines()
                for i in range(len(l)):
                    if key in l[i]:
                        out.write(l[i - 1])
                        out.write(l[i])
                    if key2 in l[i]:
                        out.write(l[i])
            out.close()

            # 删除常数变量的output节点
        def deleteConstant():
            with open('deleteConstant_output.txt', 'w') as input:
                keyword = ('op_type: "Constant"')
                f = open('output.txt', 'r')
                l = f.readlines()
                for i in range(len(l)):
                    if i<len(l)-2:
                        if keyword in l[i+2] or keyword in l[i+1] or keyword in l[i]:
                            pass
                        else:
                            input.write(l[i])
                    else:
                        input.write(l[i])
                input.close()


        # 删除非节点数据
        def drop(words: list):
            f1 = open('deleteConstant_output.txt')
            f2 = open('output_delete.txt', 'w')
            s = pandas.Series(f1.readlines())
            flag = numpy.array([True] * len(s))
            for word in words:
                flag = flag & ~s.str.contains(word)
            s = s[flag].tolist()
            f2.writelines(s)
            f1.close()
            f2.close()

            # 把多个输出流合在一起
            temp = ''
            with open('output_emerge.txt', 'w') as output:
                keyword = ('output')
                f = open('output_delete.txt', 'r')
                flag = True
                l = f.readlines()
                for i in range(len(l)):
                    if i + 1 <= len(l):
                        if keyword in l[i]:
                            if keyword in l[i + 1]:
                                flag = False
                                strinput = l[i].split('"')[1]
                                temp = (temp + strinput + ';')
                                # print(temp)
                                output.write('')
                            elif keyword not in l[i + 1] and flag == False:
                                strline = l[i].split('"')[1]
                                temp = ('output:\"' + temp + strline + '\"' + '\n')
                                output.write(temp)
                                temp = ''
                            else:
                                output.write(l[i])
                        else:
                            output.write(l[i])
                output.close()

            i = 1
            #对节点信息分类好，把每一个节点的信息分在同一行
            with open('output_outdelete.txt', 'w') as f:
                with open('output_emerge.txt', 'r') as fp:
                    for line in fp:
                        n = i % 3
                        if n != 0:
                            line = line.strip('\n')
                            f.write(line)
                        else:
                            f.write(line)
                        i = i + 1

        # 保存数据到CSV
        def xls_write_output(nosuffixname):
            book = xlwt.Workbook(encoding='utf-8', style_compression=0)
            sheet = book.add_sheet('Output', cell_overwrite_ok=True)
            sheet.write(0, 0, 'id')
            sheet.write(0, 1, 'name')
            sheet.write(0, 2, 'type')
            sheet.write(0, 3, 'output')
            n = 1
            with open('output_outdelete.txt', 'r+') as f:
                for text in f.readlines():
                    num = n
                    # name = text.split('"')[3]
                    # ran_str = ''.join(random.sample(
                    #     ['z', 'y', 'x', 'w', 'v', 'u', 't', 's', 'r', 'q', 'p', 'o', 'n', 'm', 'l', 'k', 'j', 'i', 'h',
                    #      'g', 'f', 'e', 'd', 'c', 'b', 'a'], 5))
                    name = str(n)
                    type = text.split('"')[5]
                    output = text.split('"')[1]
                    sheet.write(n, 0, num)
                    sheet.write(n, 1, name.split('"')[0])
                    sheet.write(n, 2, type.split('"')[0])
                    sheet.write(n, 3, output.split('"')[0])
                    n = n + 1
            book.save(nosuffixname + '_output.xls')

        research(nosuffixname)
        deleteConstant()
        drop(["node", "{", "}","Constant"])

        xls_write_output(nosuffixname)
        print('output success')

    # 从xls转为图结构
def xlstograph(nosuffixname):
        data_input = pandas.read_excel(nosuffixname+'_input.xls', keep_default_na=False, index_col=u'name')
        data_output = pandas.read_excel(nosuffixname+'_output.xls', keep_default_na=False, index_col=u'name')
        conter_list = Counter()

        # 根据input或output中的值定位到name
        def find_row(num_value, data):
            # data.set_index(['name'],inplace=True)
            #index 是 每一个行（节点）的名称
            for indexs in data.index:
                for i in range(len(data.loc[indexs].values)):
                    #num_value 为输入输出值
                    if (str(data.loc[indexs].values[i]) == num_value):
                        row = indexs
                        # print(row)
                        return (row)

        #根据name中的值定位到type，为节点描述category
        def find_row_type(name,data):
            type_category = {'Add':1,'Relu':2,'MaxPool':3,'Conv':4,'Reshape':5,
                             'MatMul':6,'ImageScaler':7,'BatchNormalization':8,
                             'LeakyRelu':9,'Sub':10,'Div':11,'Dropout':12,'Concat':14,
                             'AveragePool':15}
            for indexs in data.index:
                if(indexs == name):
                        # print(data.loc[indexs].values[i])
                        type_name = str(data.loc[indexs].values[1])
                        return type_category.get(type_name)
        # 双遍历找到input和output相同的值
        ori_list = {}
        set_list = []
        for x in data_input['input']:
            strlist_x = x.split(';')
            for value_x in strlist_x:
                for y in data_output['output']:
                    strlist_y = y.split(';')
                    for value_y in strlist_y:
                        # print(value_x, value_y)
                        if value_y == value_x and value_x != '' and value_y != '':
                            #从input 定位到name
                            row_x = find_row(x, data_input)
                            #从output 定位到name
                            row_y = find_row(y, data_output)
                            # print(row_y,row_x)
                            set_list += [row_y, row_x]
                            if (row_y, row_x) not in ori_list:
                                ori_list[(row_y, row_x)] = 1
                            else:
                                ori_list[(row_y, row_x)] += 1
        ori_list = sorted(ori_list.items(), key=lambda x: x[1], reverse=True)
        data_list = []
        label_x = [-10 + i * 1 for i in range(0, 49)]
        for k in ori_list:
            data_list.append(list(k[0]) + [k[1]])
        # print(data_list)
        set_list = list(set(set_list))
        add_nodes = []
        add_nodes_id = []
        we = []

        def main():
            G = nx.Graph()
            # 添加节点
            for word in data_list[:]:
                # print(word)
                for w in word[:2]:
                    if w not in add_nodes:
                        G.add_node(str(w),category=find_row_type(w,data_output),radius=1)
                        add_nodes_id.append(find_row_type(w,data_output))
                        add_nodes.append(w)
            print(G.nodes)
            print(add_nodes_id)
            # 添加边
            for item in data_list[:]:
                G.add_edge(str(item[0]), str(item[1]))
                we.append(item[2])
            # print(G.edges)
            pos = nx.spring_layout(G)
            nx.draw(G, with_labels=True, node_color='y', font_size=10, arrows=True)
            # plt.savefig('tu.pdf')
            # plt.show()
            nx.draw_networkx(G)
            # print(nx.adjacency_matrix(G).todense())
            print('to graph success')

            ## 打包 ##
            # nx.set_node_attributes(G,'pos',pos)
            nx.write_gpickle(G,nosuffixname+"_Gpi")
            print('to GPI success')
        main()

if __name__ == '__main__':
    # for i in range(1,len(sys.argv[1])):
    #     filepath = sys.argv[i]
    # filepath = "D:\model\Densenet121.onnx"

    ##########################################idea调用
    filepath = ""
    for i in range(1,len(sys.argv)):
        filepath = filepath + (sys.argv[i])
    ###########################################
    # print(filepath)
    # print(type(filepath))
    nosuffixname = filepath[0:-5]
    print(nosuffixname)
    inload(nosuffixname)
    toxls_input(nosuffixname)
    print('input success')
    toxls_output(nosuffixname)
    xlstograph(nosuffixname)
